import torch
import torch.nn as nn
import torch.nn.functional as F


class lossAV(nn.Module):

    def __init__(self):
        super(lossAV, self).__init__()
        self.criterion = nn.CrossEntropyLoss(reduction='none')
        self.FC = nn.Linear(256, 2)

    def forward(self, x, labels=None, masks=None):
        x = x.squeeze(1)
        x = self.FC(x)
        if labels is None:
            predScore = x[:, 1]
            predScore = predScore.t()
            predScore = predScore.view(-1).detach().cpu().numpy()
            return predScore
        else:
            sumloss = torch.sum(self.criterion(x, labels) * masks)
            # nloss = torch.sum(nloss) / (num_valid + 1e-6)  # 避免除 0
            predScore = F.softmax(x, dim=-1)
            # predLabel = torch.argmax(predScore, dim=-1) # 改进版 🔥
            predLabel = torch.round(predScore)[:, 1] # 原版
            correctNum = ((predLabel == labels) * masks).sum().float()
            return sumloss, predScore, predLabel, correctNum


class lossA(nn.Module):

    def __init__(self):
        super(lossA, self).__init__()
        self.criterion = nn.CrossEntropyLoss(reduction='none')
        self.FC = nn.Linear(128, 2)

    def forward(self, x, labels, masks=None):
        x = x.squeeze(1)
        x = self.FC(x)
        sumloss = torch.sum(self.criterion(x, labels) * masks)
        return sumloss


class lossV(nn.Module):

    def __init__(self):
        super(lossV, self).__init__()
        self.criterion = nn.CrossEntropyLoss(reduction='none')
        self.FC = nn.Linear(128, 2)

    def forward(self, x, labels, masks=None):
        x = x.squeeze(1)
        x = self.FC(x)
        sumloss = torch.sum(self.criterion(x, labels) * masks)
        return sumloss


class lossTime(nn.Module):

    def __init__(self):
        super(lossTime, self).__init__()
        self.criterion = nn.CrossEntropyLoss(reduction='none')
        self.FC = nn.Linear(128, 2)

    def forward(self, x, labels, masks=None):
        x = x.squeeze(1)
        x = self.FC(x)
        sumloss = torch.sum(self.criterion(x, labels) * masks)
        return sumloss

class lossSub(nn.Module):

    def __init__(self):
        super(lossSub, self).__init__()
        self.criterion = nn.CrossEntropyLoss(reduction='none')
        self.FC = nn.Linear(128, 2)

    def forward(self, x, labels, masks=None):
        x = x.squeeze(1)
        x = self.FC(x)
        sumloss = torch.sum(self.criterion(x, labels) * masks)
        return sumloss

class lossWin(nn.Module):

    def __init__(self):
        super(lossWin, self).__init__()
        self.criterion = nn.CrossEntropyLoss(reduction='none')
        self.FC = nn.Linear(128, 2)

    def forward(self, x, labels, masks=None):
        x = x.squeeze(1)
        x = self.FC(x)
        sumloss = torch.sum(self.criterion(x, labels) * masks)
        return sumloss